package util

import java.time.{Instant, LocalDateTime, ZoneId}
import java.util

import bean.DataPoint
import com.cloudera.sparkts.models.ARIMAModel
import org.apache.spark.mllib.linalg.{DenseVector, Vector}

import scala.collection.mutable.ArrayBuffer

/**
  * @author :chensi
  */
object TsUtils {

  def linearForeast(line: RegressionLine, forecastDays: Int, index: Int, preserveDays: Int) = {
    var r: Array[Double] = null
    if (preserveDays > forecastDays) {
      r = Array.ofDim[Double](forecastDays)
    } else {
      r = Array.ofDim[Double](preserveDays)
    }
    line.getA1
    var begin = index
    var rIndex = 0
    for (i <- 0 until forecastDays) {
      val forecast: Double = begin * line.getA1 + line.getA0
      if (i >= forecastDays - preserveDays) {
        r(rIndex) = forecast
        rIndex = rIndex + 1
      }
      begin = begin + 1
    }
    r
  }

  /**
    * Use ARIMA model forcast the value and extract the forecast value
    *
    * @param arimaModel
    * @param vector
    * @param forecastDays
    * @param preserveDays
    * @return
    */
  def arimaForecast(arimaModel: ARIMAModel, vector: Vector, forecastDays: Int, preserveDays: Int) = {
    try {
      var r: Array[Double] = null
      if (preserveDays > forecastDays) {
        r = Array.ofDim[Double](forecastDays)
      } else {
        r = Array.ofDim[Double](preserveDays)
      }
      val allVector = arimaModel.forecast(vector, forecastDays).toArray
      Array.copy(allVector, allVector.size - preserveDays, r, 0, preserveDays)
      r
    } catch {
      case e: Exception => null
    }
  }

  /**
    * To generate day array
    *
    * @param data
    * @return
    */
  def generateDayArr(data: util.List[DataPoint]) = {
    if (data == null || data.size() == 0) {
      null
    } else {
      val firstDay = data.get(0).getTs
      val lastDay = data.get(data.size() - 1).getTs
      val arr = Array.fill((lastDay - firstDay + 1).toInt) {
        Double.NaN
      }
      for (i <- 0 until data.size()) {
        arr(data.get(i).getTs.toInt - firstDay.toInt) = data.get(i).getValue
      }
      arr
    }
  }

  /**
    *
    * @param data 按时间排好序的DataPoint集合
    * @return 分钟级别的向量 如果参数为null或大小为0则返回null
    */
  def generateMinuteArr(data: util.List[DataPoint]) = {
    if (data == null || data.size() == 0) {
      null
    } else {
      var initIndex = -1
      //存储一分钟的DataPoint value
      val tmpArr = new ArrayBuffer[Double](60)
      //存储一天的DataPoint value
      val dayArr = Array.fill(24 * 60) {
        Double.NaN
      }
      for (i <- 0 until data.size()) {
        val point = data.get(i)
        val time = LocalDateTime.ofInstant(Instant.ofEpochMilli(point.getTs), ZoneId.of("Asia/Shanghai"))
        val index = time.getHour * 60 + time.getMinute
        if (initIndex == -1) {
          initIndex = index;
          tmpArr.+=(point.getValue)
        } else if (initIndex == index) {
          tmpArr.+=(point.getValue)
        } else {
          dayArr(index) = tmpArr.sum / tmpArr.length
          initIndex = index
          tmpArr.clear()
        }
      }
      //移除dayArr首尾的NaN
      var prefix: Int = -1
      var suffix: Int = -1
      for (i <- 0 until dayArr.length if prefix == -1) {
        if (!(dayArr(i) equals Double.NaN)) {
          prefix = i
        }
      }
      for (i <- (0 until (dayArr.length)).reverse if suffix == -1) {
        if (!(dayArr(i) equals Double.NaN)) {
          suffix = i
        }
      }
      val validArr = Array.ofDim[Double](suffix - prefix + 1)
      Array.copy(dayArr, prefix, validArr, 0, validArr.length)
      (validArr, suffix)
    }
  }

  /**
    * 将数组线性填充
    *
    * @param values
    * @return
    */
  def fillLinear(values: Vector): DenseVector = {
    val result = values.copy.toArray
    var i = 1
    while (i < result.length - 1) {
      val rangeStart = i
      while (i < result.length - 1 && result(i).isNaN) {
        i += 1
      }
      val before = result(rangeStart - 1)
      val after = result(i)
      if (i != rangeStart && !before.isNaN && !after.isNaN) {
        val increment = (after - before) / (i - (rangeStart - 1))
        for (j <- rangeStart until i) {
          result(j) = result(j - 1) + increment
        }
      }
      i += 1
    }
    new DenseVector(result)
  }

}
